# -*- coding: utf-8 -*-

import paramiko
import logging
import time
logging.basicConfig(level=logging.DEBUG,
                    format='[%(asctime)s] %(levelname)s %(module)s:%(lineno)d %(message)s')
LOG = logging.getLogger(__name__)

SHELL_ENDSWITH = ']# '
SHELL_ENDSWITH = '\x1b[34m->\x1b[01;34m %\x1b[00m \x1b[K\x1b[?1h\x1b=\x1b[?2004h'

RETRY_RECEIVE = 3
WAIT_RECEIVE = 1
WAIT_LOGIN = 3
WAIT_COMMAND = 1


def receive(channel, retry=RETRY_RECEIVE):
    data = ''
    if not channel.exit_status_ready():
        while retry > 0:
            time.sleep(WAIT_RECEIVE)
            if channel.recv_ready():
                retry = RETRY_RECEIVE
                _data = channel.recv(1024)
                # LOG.debug('get data: %s...', _data[:100])
                # LOG.debug('get data: %s', _data)
                data += _data
                if _data.endswith(SHELL_ENDSWITH):
                    break
            else:
                retry -= 1
                LOG.debug('channel is not receive ready')
    else:
        LOG.warn('channel exit')
    data = '\r\n'.join(data.split('\r\n')[1:-1])
    return data


def flush(channel):
    time.sleep(WAIT_COMMAND)
    return receive(channel)


def remote_execute(channel, command, show=False):
    LOG.info('COMMAND: %r', command)
    channel.send(command + '\n')
    content = flush(channel)
    if show:
        for line in content.splitlines():
            LOG.debug('RESPONSE: %s', line)
    return content


def start(connections):
    try:
        c = paramiko.SSHClient()
        c.set_missing_host_key_policy(paramiko.AutoAddPolicy())

        if isinstance(connections, dict):
            connections = [connections]

        # hostname, port, username, password, sudo
        c.connect(**connections[0])
        channel = c.invoke_shell()
        flush(channel)
        remote_execute(channel, 'bash --noprofile --norc', show=True)

        for connection in connections[1:]:
            hostname = connection.get('hostname')
            port = connection.get('port', 22)

            username = connection.get('username')
            as_root = connection.get('as_root', False)

            if username:
                command = 'ssh %s@%s -p %d' % (username, hostname, port)
            else:
                command = 'ssh %s -p %d' % (hostname, port)

            if as_root:
                command = 'sudo ' + command

            remote_execute(channel, command, show=False)
            time.sleep(WAIT_LOGIN)
            remote_execute(channel, 'bash --noprofile --norc', show=True)

        return channel

    except Exception as e:
        logging.exception("Connection Failed: %s", e)


def main():
    chan = start([
        {'hostname': 'localhost', 'username': 'catroll', 'password': '532084'},
    ])
    remote_execute(chan, 'LANG=en free', show=True)


if __name__ == '__main__':
    main()
